#pragma once

#include "stair_ctrl/lstm.hpp"
#include "stair_ctrl/interface.hpp"
#include "stair_ctrl/step_plan_interface.hpp"
#include "stair_ctrl/NN_log.hpp"

#include <Eigen/Core>

#include <iostream>
#include <cmath>


class NNController
{
private:
    LstmPolicy lstm_policy;

    int ob_dim, act_dim;

    Eigen::VectorXf input, output, input_raw, joint_last;
    Eigen::VectorXf ob_mean, ob_std, act_mean, act_std;
    Eigen::VectorXf end_effector_offset;
    
    double vx_max = 1.0;
    double vy_max = 0.4;
    double omega_max = 1.0;

    Eigen::MatrixXf NN_log;
    int NN_log_index, NN_log_buffer;

public:
    NNController(std::string path = "")
    {
        lstm_policy = LstmPolicy(path);
        lstm_policy.reset();

        ob_dim = 62;
        act_dim = 12;

        input.setZero(ob_dim);
        ob_mean.setZero(ob_dim);
        ob_std.setZero(ob_dim);
        output.setZero(act_dim);
        act_mean.setZero(act_dim);
        act_std.setZero(act_dim);
        input_raw.setZero(ob_dim);
        joint_last.setZero(12);

        end_effector_offset.setZero(12);

        act_mean << 0.0, 0.8, -1.75,
                    0.0, 0.8, -1.75,
                    0.0, 0.8, -1.75,
                    0.0, 0.8, -1.75;
        act_std << 0.4, 1.2, 0.9,
                    0.4, 1.2, 0.9,
                    0.4, 1.2, 0.9,
                    0.4, 1.2, 0.9;
        end_effector_offset << 0.1881, -0.12675, -0.3,
                                0.1881, 0.12675, -0.3,
                                -0.1881, -0.12675, -0.3,
                                -0.1881, 0.12675, -0.3;
        ob_mean << 0.0, 0.0, 0.0,                           // cmd
                    0.0, 0.0,                               // phase
                    act_mean,                               // jpos
                    Eigen::VectorXf::Constant(12, 0.0),     // jvel
                    0.0, 0.0, 1.0,                          // r_g in body
                    0.0, 0.0, 0.0,                          // body ang vel
                    end_effector_offset,
                    end_effector_offset,
                    0.0, 0.0, 1.0;                          // r_g in body des
        ob_std << 1.0/1.5, 1.0/0.8, 1.0/1.0,
                    1.0, 1.0,
                    act_std.cwiseInverse(),
                    1.0/24.0, 1.0/40.0, 1.0/40.0,
                    1.0/24.0, 1.0/40.0, 1.0/40.0,
                    1.0/24.0, 1.0/40.0, 1.0/40.0,
                    1.0/24.0, 1.0/40.0, 1.0/40.0,
                    1.0/0.3, 1.0/0.3, 1.0/0.3,              // r_g in body
                    1.0/3.0, 1.0/3.0, 1.0/3.0,              // body ang vel
                    1.0/0.3, 1.0/0.15, 1.0/0.15,
                    1.0/0.3, 1.0/0.15, 1.0/0.15,
                    1.0/0.3, 1.0/0.15, 1.0/0.15,
                    1.0/0.3, 1.0/0.15, 1.0/0.15,
                    1.0/0.3, 1.0/0.15, 1.0/0.15,
                    1.0/0.3, 1.0/0.15, 1.0/0.15,
                    1.0/0.3, 1.0/0.15, 1.0/0.15,
                    1.0/0.3, 1.0/0.15, 1.0/0.15,
                    1.0/0.3, 1.0/0.3, 1.0/0.3;
        
        NN_log_index = 0;
        NN_log_buffer = 100;
        NN_log.setZero(NN_log_buffer, ob_dim+12);
    }

    void update(robotInterface<float> &robot_interface, cmdInterface<float> &cmd_interface, stepPlanInterface<float> &step_plan_interface)
    {
        // cmd
        input_raw[0] = cmd_interface.cmd[0];
        input_raw[1] = cmd_interface.cmd[1];
        input_raw[2] = cmd_interface.cmd[2];

        // phase
        input_raw.segment<2>(3) << std::sin(step_plan_interface.phase_scalar), std::cos(step_plan_interface.phase_scalar);

        // joint pos and vel
        input_raw.segment<12>(5) << robot_interface.joint_pos;
        input_raw.segment<12>(17) << robot_interface.joint_vel_filtered;

        // imu
        auto mat = robot_interface.getRotMat(robot_interface.body_quat);
        input_raw.segment<3>(29) << mat(0, 2), mat(1, 2), mat(2, 2);
        input_raw.segment<3>(32) << robot_interface.body_ang_vel;

        // foot and body pos des
        for (int i = 0; i < 4; ++i)
        {
            if (step_plan_interface.motion_plan_phase_vec[i] < M_PI)
            {
                input_raw.segment<3>(35+3*i) << step_plan_interface.end_effector_ref.segment<3>(3*i);
                input_raw.segment<3>(47+3*i) << step_plan_interface.end_effector_ref_prev.segment<3>(3*i);
            }
            else
            {
                input_raw.segment<3>(35+3*i) << step_plan_interface.end_effector_ref_prev.segment<3>(3*i);
                input_raw.segment<3>(47+3*i) << step_plan_interface.end_effector_ref.segment<3>(3*i);
            }
        }
        input_raw.tail<3>() << step_plan_interface.rot_des;

        // update policy
        input = (input_raw - ob_mean).cwiseProduct(ob_std);
        lstm_policy.input = input;
        lstm_policy.predict();
        output = lstm_policy.output.head<12>();
        output = output.cwiseProduct(act_std) + act_mean;

        // set command
        robot_interface.joint_pos_des << output;
    }

    void updateLog(std::string NN_log_file)
    {
        NN_log.row(NN_log_index).head(ob_dim) = input_raw;
        NN_log.row(NN_log_index).segment<12>(ob_dim) = output;

        // record NN log
        ++NN_log_index;
        if (NN_log_index > NN_log_buffer - 1) {
            NN_log_index = 0;
            auto log_status = WriteLog(NN_log_file, NN_log);
            if (log_status == 0) {
                std::cout << "[SUCCESS] Save log at " << NN_log_file << std::endl;
            } else {
                std::cout << "[FAIL] Attempt to save log at " << NN_log_file << std::endl;
            }
            NN_log.setZero(NN_log_buffer, ob_dim+12);
        }
    }

    Eigen::VectorXf getNNInput()
    {
        return input;
    }

    Eigen::VectorXf getNNInputRaw()
    {
        return input_raw;
    }

    Eigen::VectorXf getNNOutput()
    {
        return output;
    }

};